#include "Mesh.h"

using namespace OBR;

/* ------------------------ initialize static members ----------------------- */
size_t Mesh::gas_total_built_sbt = 0;

Mesh::Mesh(std::string cuda_source_path)
    : ASNode(OBR::ASNodeType::GAS), VirtualProgramEntry(cuda_source_path)
{
}

Mesh::~Mesh()
{
}

void Mesh::buildAccelStructure(OptixDeviceContext &ctx, CUstream &stream)
{
    if (!this->built)
    {
        OptixAccelBuildOptions opt = this->getASBuildOptions();
        OptixBuildInput input = this->getASBuildInput();

        /* ----------------------------- calc sizes ----------------------------- */
        OptixAccelBufferSizes gas_buffer_sizes;
        OPTIX_CHECK(optixAccelComputeMemoryUsage(ctx, &opt, &input,
                                                 1, // num_build_inputs
                                                 &gas_buffer_sizes));

        CudaBuffer d_temp_buffer(gas_buffer_sizes.tempSizeInBytes);

        CUdeviceptr d_buffer_temp_output_gas_and_compacted_size; // non-compacted output
        size_t compactedSizeOffset =
            (gas_buffer_sizes.outputSizeInBytes + 8 - 1) / 8 * 8; // round up
        CUDA_CHECK(
            cudaMalloc(reinterpret_cast<void **>(&d_buffer_temp_output_gas_and_compacted_size),
                       compactedSizeOffset + 8));

        OptixAccelEmitDesc emitProperty = {};
        emitProperty.type = OPTIX_PROPERTY_TYPE_COMPACTED_SIZE;
        emitProperty.result = (CUdeviceptr)((char *)d_buffer_temp_output_gas_and_compacted_size +
                                            compactedSizeOffset);

        /* -------------------------------- build ------------------------------- */
        OPTIX_CHECK(optixAccelBuild(ctx, stream, &opt, &input,
                                    1, // num build inputs
                                    d_temp_buffer.d_pointer(), gas_buffer_sizes.tempSizeInBytes,
                                    d_buffer_temp_output_gas_and_compacted_size,
                                    gas_buffer_sizes.outputSizeInBytes, &this->handle,
                                    &emitProperty, // emitted property list
                                    1              // num emitted properties
                                    ));

        /* ------------------------------- compact ------------------------------ */
        size_t compacted_gas_size;
        CUDA_CHECK(cudaMemcpy(&compacted_gas_size, (void *)emitProperty.result, sizeof(size_t),
                              cudaMemcpyDeviceToHost));

        if (compacted_gas_size < gas_buffer_sizes.outputSizeInBytes)
        {
            this->d_gas_output.resize(compacted_gas_size);

            // use handle as input and output
            OPTIX_CHECK(optixAccelCompact(ctx, 0, this->handle, this->d_gas_output.d_pointer(),
                                          compacted_gas_size, &this->handle));

            CUDA_CHECK(cudaFree((void *)d_buffer_temp_output_gas_and_compacted_size));
        }
        else
        {
            this->d_gas_output.resize(compactedSizeOffset + 8);
            CUDA_CHECK(
                cudaMemcpy(reinterpret_cast<void *>(this->d_gas_output.d_pointer()),
                           reinterpret_cast<void *>(d_buffer_temp_output_gas_and_compacted_size),
                           compactedSizeOffset + 8, cudaMemcpyDeviceToDevice));
        }

        /* ---------------------------- finishing up ---------------------------- */
        built = true;
        this->gas_sbt_base = Mesh::gas_total_built_sbt;
        Mesh::gas_total_built_sbt += this->ref_cnt;
    }
}